import pandas as pd
import requests
import time
import random
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

def lookup_pubchem_cid(name, max_retries=3):
    """
    Look up PubChem Compound CID for a given chemical name using PUG REST API.
    
    Args:
        name (str): Chemical name to search
        max_retries (int): Maximum number of retry attempts
    
    Returns:
        str or None: Compound CID if found, None otherwise
    """
    # Base URLs for PubChem API
    search_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name"
    
    # Attempts to find the best match for the chemical name
    for attempt in range(max_retries):
        try:
            # Encode the name to handle special characters
            encoded_name = requests.utils.quote(name)
            
            # First, try to get the CID directly
            response = requests.get(
                f"{search_url}/{encoded_name}/CID/TXT", 
                timeout=10
            )
            
            # If direct lookup fails, try a more comprehensive search
            if response.status_code != 200 or not response.text.strip().isdigit():
                # Use the identifier exchange to find the CID
                search_response = requests.get(
                    f"{search_url}/{encoded_name}/cids/TXT",
                    timeout=10
                )
                
                if search_response.status_code == 200:
                    # Split and take the first CID
                    cids = search_response.text.strip().split()
                    if cids:
                        return cids[0]
            else:
                # Direct CID lookup successful
                return response.text.strip()
            
            # Exponential backoff with jitter if lookup fails
            backoff_time = (0.5 * (2 ** attempt)) + (random.random() * 0.1)
            time.sleep(backoff_time)
        
        except requests.exceptions.RequestException as e:
            # Exponential backoff with jitter
            backoff_time = (0.5 * (2 ** attempt)) + (random.random() * 0.1)
            time.sleep(backoff_time)
        except Exception as e:
            time.sleep(0.5 * (2 ** attempt))
    
    return None

def get_pubchem_smiles(cid, max_retries=3):
    """
    Get SMILES notation for a compound using its PubChem CID.
    
    Args:
        cid (str): PubChem Compound ID
        max_retries (int): Maximum number of retry attempts
    
    Returns:
        str or None: SMILES notation if found, None otherwise
    """
    # Base URL for PubChem property API
    property_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid"
    
    for attempt in range(max_retries):
        try:
            # Request SMILES notation for the CID
            response = requests.get(
                f"{property_url}/{cid}/property/CanonicalSMILES/TXT",
                timeout=10
            )
            
            if response.status_code == 200:
                smiles = response.text.strip()
                return smiles
            
            # Exponential backoff with jitter if lookup fails
            backoff_time = (0.5 * (2 ** attempt)) + (random.random() * 0.1)
            time.sleep(backoff_time)
            
        except requests.exceptions.RequestException as e:
            # Exponential backoff with jitter
            backoff_time = (0.5 * (2 ** attempt)) + (random.random() * 0.1)
            time.sleep(backoff_time)
        except Exception as e:
            time.sleep(0.5 * (2 ** attempt))
    
    return None

def lookup_pubchem_data(name):
    """
    Combined function to look up both CID and SMILES for a given compound name.
    
    Args:
        name (str): Chemical name to search
    
    Returns:
        tuple: (cid, smiles) pair, either may be None if not found
    """
    cid = lookup_pubchem_cid(name)
    smiles = None
    
    if cid:
        smiles = get_pubchem_smiles(cid)
    
    return (cid, smiles)

def batch_pubchem_lookup(df, name_col='NameEN', cid_col='CID', smiles_col='Smiles', batch_size=100):
    """
    Batch lookup PubChem CIDs and SMILES for rows starting with 'Miss_CID'
    
    Args:
        df (pd.DataFrame): Input DataFrame
        name_col (str): Column containing chemical names
        cid_col (str): Column to update with PubChem CIDs
        smiles_col (str): Column to update with SMILES notations
        batch_size (int): Number of rows to process in each batch
    
    Returns:
        pd.DataFrame: Updated DataFrame with lookup statistics
    """
    # Create a copy to avoid modifying the original
    updated_df = df.copy()
    
    # Add SMILES column if it doesn't exist
    if smiles_col not in updated_df.columns:
        updated_df[smiles_col] = None
    
    # Filter rows that need CID lookup
    mask = updated_df[cid_col].str.startswith('Miss_CID')
    lookup_df = updated_df[mask]
    
    # Tracking variables
    total_rows = len(lookup_df)
    successful_cid_lookups = 0
    failed_cid_lookups = 0
    successful_smiles_lookups = 0
    failed_smiles_lookups = 0
    
    # Progress bar for the entire process
    with tqdm(total=total_rows, desc="PubChem Lookup", 
              bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]') as pbar:
        # Process in batches to manage memory and API load
        for start in range(0, total_rows, batch_size):
            batch = lookup_df.iloc[start:start+batch_size]
            
            # Use ThreadPoolExecutor for concurrent API requests
            with ThreadPoolExecutor(max_workers=min(10, batch_size)) as executor:
                # Create futures for each combined lookup (CID + SMILES)
                futures = {
                    executor.submit(lookup_pubchem_data, row[name_col]): index 
                    for index, row in batch.iterrows()
                }
                
                # Process results as they complete
                for future in as_completed(futures):
                    index = futures[future]
                    try:
                        cid, smiles = future.result()
                        
                        # Update CID information
                        if cid:
                            updated_df.at[index, cid_col] = cid
                            successful_cid_lookups += 1
                        else:
                            failed_cid_lookups += 1
                            
                        # Update SMILES information
                        if smiles:
                            updated_df.at[index, smiles_col] = smiles
                            successful_smiles_lookups += 1
                        else:
                            failed_smiles_lookups += 1
                        
                    except Exception as e:
                        failed_cid_lookups += 1
                        failed_smiles_lookups += 1
                    
                    # Update progress bar
                    pbar.update(1)
            
            # Add a pause between batches to avoid overwhelming the API
            time.sleep(2)
    
    # Print final lookup statistics
    print("\nLookup Statistics:")
    print(f"Total Rows Processed: {total_rows}")
    print(f"Successful CID Lookups: {successful_cid_lookups}")
    print(f"Failed CID Lookups: {failed_cid_lookups}")
    print(f"CID Success Rate: {successful_cid_lookups/total_rows*100:.2f}%")
    print(f"Successful SMILES Lookups: {successful_smiles_lookups}")
    print(f"Failed SMILES Lookups: {failed_smiles_lookups}")
    if successful_cid_lookups > 0:
        print(f"SMILES Success Rate: {successful_smiles_lookups/successful_cid_lookups*100:.2f}%")
    
    return updated_df

def main():
    # Read the Excel file
    input_file = 'DB_API_temp.xlsx'
    df = pd.read_excel(input_file)
    
    # Perform CID and SMILES lookup
    try:
        updated_df = batch_pubchem_lookup(df)
        
        # Save updated DataFrame
        output_file = 'DB_API_temp_Updated.xlsx'
        updated_df.to_excel(output_file, index=False)
        
        print(f"\nCID and SMILES lookup completed. Results saved to {output_file}")
    
    except Exception as e:
        print(f"Batch lookup failed: {e}")

if __name__ == '__main__':
    main()

# Additional considerations:
# 1. Test with a small subset of data first
# 2. Manual verification of critical compounds recommended
# 3. Consider adding an option to look up SMILES for compounds that already have CIDs but are missing SMILES